#!/usr/bin/env python3
"""
机器学习引擎 - 验证脚本

本脚本用于验证新的ml_engine目录结构是否正常工作。

使用方法：
    python verify_ml_engine.py
"""

import os
import sys
import traceback

def test_imports():
    """测试模块导入"""
    print("="*60)
    print("测试模块导入")
    print("="*60)
    
    try:
        # 测试基础导入
        print("1. 测试基础导入...")
        import numpy as np
        import pandas as pd
        import sklearn
        print("   ✓ 基础依赖导入成功")
        
        # 测试ML引擎导入
        print("2. 测试ML引擎导入...")
        from ml_engine import (
            EnhancedProblemDataset,
            EnhancedMLProblemClassifier,
            MLPredictor,
            MLModelManager,
            config
        )
        print("   ✓ ML引擎核心模块导入成功")
        
        # 测试便捷函数导入
        print("3. 测试便捷函数导入...")
        from ml_engine import train_model, predict, load_model, list_models
        print("   ✓ 便捷函数导入成功")
        
        return True
        
    except Exception as e:
        print(f"   ✗ 导入失败: {e}")
        traceback.print_exc()
        return False

def test_dataset():
    """测试数据集功能"""
    print("\n" + "="*60)
    print("测试数据集功能")
    print("="*60)
    
    try:
        from ml_engine.dataset import EnhancedProblemDataset
        
        # 创建数据集
        print("1. 创建增强数据集...")
        dataset = EnhancedProblemDataset()
        print("   ✓ 数据集创建成功")
        
        # 获取训练数据
        print("2. 获取训练数据...")
        training_data = dataset.get_enhanced_training_data()
        print(f"   ✓ 获取到 {len(training_data)} 个训练样本")
        
        # 获取统计信息
        print("3. 获取统计信息...")
        stats = dataset.get_dataset_statistics()
        print(f"   ✓ 数据集统计: {stats['total_samples']} 个样本, {stats['unique_labels']} 个标签")
        
        # 保存数据集
        print("4. 保存数据集...")
        dataset.save_enhanced_dataset("test_dataset.json")
        print("   ✓ 数据集保存成功")
        
        return True
        
    except Exception as e:
        print(f"   ✗ 数据集测试失败: {e}")
        traceback.print_exc()
        return False

def test_trainer():
    """测试训练器功能"""
    print("\n" + "="*60)
    print("测试训练器功能")
    print("="*60)
    
    try:
        from ml_engine.trainer import EnhancedMLProblemClassifier
        
        # 创建分类器
        print("1. 创建分类器...")
        classifier = EnhancedMLProblemClassifier("naive_bayes")
        print("   ✓ 分类器创建成功")
        
        # 准备测试数据
        print("2. 准备测试数据...")
        test_texts = [
            "计算两个整数的和",
            "找到数组中的最大值",
            "反转字符串",
            "计算阶乘",
            "判断质数"
        ]
        test_labels = [
            "addition",
            "array_max",
            "string_reverse",
            "factorial",
            "prime_check"
        ]
        print(f"   ✓ 准备了 {len(test_texts)} 个测试样本")
        
        # 训练模型
        print("3. 训练模型...")
        results = classifier.train(test_texts, test_labels)
        print(f"   ✓ 模型训练完成，准确率: {results['accuracy']:.4f}")
        
        # 测试预测
        print("4. 测试预测...")
        prediction, confidence = classifier.predict("计算三个数的和")
        print(f"   ✓ 预测结果: {prediction} (置信度: {confidence:.3f})")
        
        return True
        
    except Exception as e:
        print(f"   ✗ 训练器测试失败: {e}")
        traceback.print_exc()
        return False

def test_predictor():
    """测试预测器功能"""
    print("\n" + "="*60)
    print("测试预测器功能")
    print("="*60)
    
    try:
        from ml_engine.predictor import MLPredictor, MLModelManager
        
        # 创建预测器
        print("1. 创建预测器...")
        predictor = MLPredictor()
        print("   ✓ 预测器创建成功")
        
        # 测试模型管理器
        print("2. 测试模型管理器...")
        manager = MLModelManager()
        models = manager.list_models()
        print(f"   ✓ 找到 {len(models)} 个模型文件")
        
        # 如果有模型，测试预测
        if models:
            print("3. 测试模型预测...")
            model_file = models[0]['filename']
            if predictor.load_model(model_file):
                prediction, confidence = predictor.predict("计算两个数的和")
                print(f"   ✓ 预测结果: {prediction} (置信度: {confidence:.3f})")
            else:
                print("   ⚠ 模型加载失败，跳过预测测试")
        else:
            print("   ⚠ 没有找到模型文件，跳过预测测试")
        
        return True
        
    except Exception as e:
        print(f"   ✗ 预测器测试失败: {e}")
        traceback.print_exc()
        return False

def test_config():
    """测试配置功能"""
    print("\n" + "="*60)
    print("测试配置功能")
    print("="*60)
    
    try:
        from ml_engine.config.settings import MLConfig
        
        # 创建配置管理器
        print("1. 创建配置管理器...")
        ml_config = MLConfig()
        print("   ✓ 配置管理器创建成功")
        
        # 获取配置
        print("2. 获取配置信息...")
        config_dict = ml_config.get_config_dict()
        print(f"   ✓ 获取到 {len(config_dict)} 个配置节")
        
        # 验证配置
        print("3. 验证配置...")
        validation = ml_config.validate_config()
        print(f"   ✓ 配置验证结果: {'有效' if validation['valid'] else '无效'}")
        
        # 更新配置
        print("4. 更新配置...")
        ml_config.update_config('training', test_size=0.3)
        print("   ✓ 配置更新成功")
        
        return True
        
    except Exception as e:
        print(f"   ✗ 配置测试失败: {e}")
        traceback.print_exc()
        return False

def test_directory_structure():
    """测试目录结构"""
    print("\n" + "="*60)
    print("测试目录结构")
    print("="*60)
    
    try:
        # 检查必要的目录
        required_dirs = [
            'ml_engine',
            'ml_engine/data',
            'ml_engine/data/raw',
            'ml_engine/data/processed',
            'ml_engine/data/enhanced',
            'ml_engine/models',
            'ml_engine/models/trained',
            'ml_engine/models/checkpoints',
            'ml_engine/models/logs',
            'ml_engine/config'
        ]
        
        print("1. 检查目录结构...")
        for dir_path in required_dirs:
            if os.path.exists(dir_path):
                print(f"   ✓ {dir_path}")
            else:
                print(f"   ✗ {dir_path} (不存在)")
                return False
        
        # 检查必要的文件
        required_files = [
            'ml_engine/__init__.py',
            'ml_engine/dataset.py',
            'ml_engine/trainer.py',
            'ml_engine/predictor.py',
            'ml_engine/train.py',
            'ml_engine/config/settings.py',
            'ml_engine/README.md'
        ]
        
        print("2. 检查文件结构...")
        for file_path in required_files:
            if os.path.exists(file_path):
                print(f"   ✓ {file_path}")
            else:
                print(f"   ✗ {file_path} (不存在)")
                return False
        
        return True
        
    except Exception as e:
        print(f"   ✗ 目录结构测试失败: {e}")
        traceback.print_exc()
        return False

def main():
    """主函数"""
    print("机器学习引擎验证脚本")
    print("="*60)
    
    success_count = 0
    total_tests = 6
    
    # 测试1: 目录结构
    if test_directory_structure():
        success_count += 1
    
    # 测试2: 模块导入
    if test_imports():
        success_count += 1
    
    # 测试3: 数据集功能
    if test_dataset():
        success_count += 1
    
    # 测试4: 训练器功能
    if test_trainer():
        success_count += 1
    
    # 测试5: 预测器功能
    if test_predictor():
        success_count += 1
    
    # 测试6: 配置功能
    if test_config():
        success_count += 1
    
    # 显示最终结果
    print("\n" + "="*60)
    print("验证结果汇总")
    print("="*60)
    
    print(f"测试通过: {success_count}/{total_tests}")
    
    if success_count == total_tests:
        print("✅ 所有测试通过！机器学习引擎可以正常使用。")
        print("\n🚀 下一步:")
        print("1. 运行完整训练: python ml_engine/train.py --train")
        print("2. 测试模型性能: python ml_engine/train.py --test")
        print("3. 集成到主系统: 更新 web_app.py 使用ML引擎")
    else:
        print("❌ 部分测试失败，请检查错误信息并修复问题。")
    
    return success_count == total_tests

if __name__ == "__main__":
    main()
